#!/usr/bin/env python
# -*- coding:utf-8 -*-

import pandas as pd
import requests
class stock_index_data:
    '''
    东方财富网-指数行情数据
    '''
    def __init__(self):
        '''
        东方财富网-指数行情数据
        '''
        pass
    def index_code_id_map_em(self) -> dict:
        """
        东方财富-股票和市场代码
        https://quote.eastmoney.com/center/gridlist.html#hs_a_board
        :return: 股票和市场代码
        :rtype: dict
        """
        url = "https://80.push2.eastmoney.com/api/qt/clist/get"
        params = {
            "pn": "1",
            "pz": "10000",
            "po": "1",
            "np": "1",
            "ut": "bd1d9ddb04089700cf9c27f6f7426281",
            "fltt": "2",
            "invt": "2",
            "fid": "f3",
            "fs": "m:1 t:2,m:1 t:23",
            "fields": "f12",
            "_": "1623833739532",
        }
        r = requests.get(url, params=params)
        data_json = r.json()
        if not data_json["data"]["diff"]:
            return dict()
        temp_df = pd.DataFrame(data_json["data"]["diff"])
        temp_df["market_id"] = 1
        temp_df.columns = ["sh_code", "sh_id"]
        code_id_dict = dict(zip(temp_df["sh_code"], temp_df["sh_id"]))
        params = {
            "pn": "1",
            "pz": "10000",
            "po": "1",
            "np": "1",
            "ut": "bd1d9ddb04089700cf9c27f6f7426281",
            "fltt": "2",
            "invt": "2",
            "fid": "f3",
            "fs": "m:0 t:6,m:0 t:80",
            "fields": "f12",
            "_": "1623833739532",
        }
        r = requests.get(url, params=params)
        data_json = r.json()
        if not data_json["data"]["diff"]:
            return dict()
        temp_df_sz = pd.DataFrame(data_json["data"]["diff"])
        temp_df_sz["sz_id"] = 0
        code_id_dict.update(dict(zip(temp_df_sz["f12"], temp_df_sz["sz_id"])))
        params = {
            "pn": "1",
            "pz": "10000",
            "po": "1",
            "np": "1",
            "ut": "bd1d9ddb04089700cf9c27f6f7426281",
            "fltt": "2",
            "invt": "2",
            "fid": "f3",
            "fs": "m:0 t:81 s:2048",
            "fields": "f12",
            "_": "1623833739532",
        }
        r = requests.get(url, params=params)
        data_json = r.json()
        if not data_json["data"]["diff"]:
            return dict()
        temp_df_sz = pd.DataFrame(data_json["data"]["diff"])
        temp_df_sz["bj_id"] = 0
        code_id_dict.update(dict(zip(temp_df_sz["f12"], temp_df_sz["bj_id"])))
        code_id_dict = {
            key: value - 1 if value == 1 else value + 1
            for key, value in code_id_dict.items()
        }
        return code_id_dict


    def index_zh_a_hist(self,
            symbol: str = "000859",
            period: str = "daily",
            start_date: str = "19700101",
            end_date: str = "22220101",
    ) -> pd.DataFrame:
        """
        东方财富网-中国股票指数-行情数据
        https://quote.eastmoney.com/zz/2.000859.html
        :param symbol: 指数代码
        :type symbol: str
        :param period: choice of {'daily', 'weekly', 'monthly'}
        :type period: str
        :param start_date: 开始日期
        :type start_date: str
        :param end_date: 结束日期
        :type end_date: str
        :return: 行情数据
        :rtype: pandas.DataFrame
        """
        code_id_dict = self.index_code_id_map_em()
        period_dict = {"daily": "101", "weekly": "102", "monthly": "103"}
        url = "https://push2his.eastmoney.com/api/qt/stock/kline/get"
        try:
            params = {
                "secid": f"{code_id_dict[symbol]}.{symbol}",
                "ut": "7eea3edcaed734bea9cbfc24409ed989",
                "fields1": "f1,f2,f3,f4,f5,f6",
                "fields2": "f51,f52,f53,f54,f55,f56,f57,f58,f59,f60,f61",
                "klt": period_dict[period],
                "fqt": "0",
                "beg": "0",
                "end": "20500000",
                "_": "1623766962675",
            }
        except KeyError:
            params = {
                "secid": f"1.{symbol}",
                "ut": "7eea3edcaed734bea9cbfc24409ed989",
                "fields1": "f1,f2,f3,f4,f5,f6",
                "fields2": "f51,f52,f53,f54,f55,f56,f57,f58,f59,f60,f61",
                "klt": period_dict[period],
                "fqt": "0",
                "beg": "0",
                "end": "20500000",
                "_": "1623766962675",
            }
            r = requests.get(url, params=params)
            data_json = r.json()
            if data_json["data"] is None:
                params = {
                    "secid": f"0.{symbol}",
                    "ut": "7eea3edcaed734bea9cbfc24409ed989",
                    "fields1": "f1,f2,f3,f4,f5,f6",
                    "fields2": "f51,f52,f53,f54,f55,f56,f57,f58,f59,f60,f61",
                    "klt": period_dict[period],
                    "fqt": "0",
                    "beg": "0",
                    "end": "20500000",
                    "_": "1623766962675",
                }
                r = requests.get(url, params=params)
                data_json = r.json()
                if data_json["data"] is None:
                    params = {
                        "secid": f"2.{symbol}",
                        "ut": "7eea3edcaed734bea9cbfc24409ed989",
                        "fields1": "f1,f2,f3,f4,f5,f6",
                        "fields2": "f51,f52,f53,f54,f55,f56,f57,f58,f59,f60,f61",
                        "klt": period_dict[period],
                        "fqt": "0",
                        "beg": "0",
                        "end": "20500000",
                        "_": "1623766962675",
                    }
                    r = requests.get(url, params=params)
                    data_json = r.json()
                    if data_json["data"] is None:
                        params = {
                            "secid": f"47.{symbol}",
                            "ut": "7eea3edcaed734bea9cbfc24409ed989",
                            "fields1": "f1,f2,f3,f4,f5,f6",
                            "fields2": "f51,f52,f53,f54,f55,f56,f57,f58,f59,f60,f61",
                            "klt": period_dict[period],
                            "fqt": "0",
                            "beg": "0",
                            "end": "20500000",
                            "_": "1623766962675",
                        }
        r = requests.get(url, params=params)
        data_json = r.json()
        try:
            temp_df = pd.DataFrame(
                [item.split(",") for item in data_json["data"]["klines"]]
            )
        except:
            # 兼容 000859(中证国企一路一带) 和 000861(中证央企创新)
            params = {
                "secid": f"2.{symbol}",
                "ut": "7eea3edcaed734bea9cbfc24409ed989",
                "fields1": "f1,f2,f3,f4,f5,f6",
                "fields2": "f51,f52,f53,f54,f55,f56,f57,f58,f59,f60,f61",
                "klt": period_dict[period],
                "fqt": "0",
                "beg": "0",
                "end": "20500000",
                "_": "1623766962675",
            }
            r = requests.get(url, params=params)
            data_json = r.json()
            temp_df = pd.DataFrame(
                [item.split(",") for item in data_json["data"]["klines"]]
            )
        temp_df.columns = [
            "日期",
            "开盘",
            "收盘",
            "最高",
            "最低",
            "成交量",
            "成交额",
            "振幅",
            "涨跌幅",
            "涨跌额",
            "换手率",
        ]
        temp_df.index = pd.to_datetime(temp_df["日期"], errors="coerce")
        temp_df = temp_df[start_date:end_date]
        temp_df.reset_index(inplace=True, drop=True)
        temp_df["开盘"] = pd.to_numeric(temp_df["开盘"], errors="coerce")
        temp_df["收盘"] = pd.to_numeric(temp_df["收盘"], errors="coerce")
        temp_df["最高"] = pd.to_numeric(temp_df["最高"], errors="coerce")
        temp_df["最低"] = pd.to_numeric(temp_df["最低"], errors="coerce")
        temp_df["成交量"] = pd.to_numeric(temp_df["成交量"], errors="coerce")
        temp_df["成交额"] = pd.to_numeric(temp_df["成交额"], errors="coerce")
        temp_df["振幅"] = pd.to_numeric(temp_df["振幅"], errors="coerce")
        temp_df["涨跌幅"] = pd.to_numeric(temp_df["涨跌幅"], errors="coerce")
        temp_df["涨跌额"] = pd.to_numeric(temp_df["涨跌额"], errors="coerce")
        temp_df["换手率"] = pd.to_numeric(temp_df["换手率"], errors="coerce")
        return temp_df


    def index_zh_a_hist_min_em(self,
            symbol: str = "399006",
            period: str = "1",
            start_date: str = "1979-09-01 09:32:00",
            end_date: str = "2222-01-01 09:32:00",
    ) -> pd.DataFrame:
        """
        东方财富网-指数数据-每日分时行情
        https://quote.eastmoney.com/center/hszs.html
        :param symbol: 指数代码
        :type symbol: str
        :param period: choice of {'1', '5', '15', '30', '60'}
        :type period: str
        :param start_date: 开始日期
        :type start_date: str
        :param end_date: 结束日期
        :type end_date: str
        :return: 每日分时行情
        :rtype: pandas.DataFrame
        """
        code_id_dict = self.index_code_id_map_em()
        if period == "1":
            url = "https://push2his.eastmoney.com/api/qt/stock/trends2/get"
            try:
                params = {
                    "fields1": "f1,f2,f3,f4,f5,f6,f7,f8,f9,f10,f11,f12,f13",
                    "fields2": "f51,f52,f53,f54,f55,f56,f57,f58",
                    "ut": "fa5fd1943c7b386f172d6893dbfba10b",
                    "iscr": "0",
                    "ndays": "5",
                    "secid": f"{code_id_dict[symbol]}.{symbol}",
                    "_": "1623766962675",
                }
            except KeyError:
                params = {
                    "fields1": "f1,f2,f3,f4,f5,f6,f7,f8,f9,f10,f11,f12,f13",
                    "fields2": "f51,f52,f53,f54,f55,f56,f57,f58",
                    "ut": "fa5fd1943c7b386f172d6893dbfba10b",
                    "iscr": "0",
                    "ndays": "5",
                    "secid": f"1.{symbol}",
                    "_": "1623766962675",
                }
                r = requests.get(url, params=params)
                data_json = r.json()
                if data_json["data"] is None:
                    params = {
                        "fields1": "f1,f2,f3,f4,f5,f6,f7,f8,f9,f10,f11,f12,f13",
                        "fields2": "f51,f52,f53,f54,f55,f56,f57,f58",
                        "ut": "fa5fd1943c7b386f172d6893dbfba10b",
                        "iscr": "0",
                        "ndays": "5",
                        "secid": f"0.{symbol}",
                        "_": "1623766962675",
                    }
                    r = requests.get(url, params=params)
                    data_json = r.json()
                    if data_json["data"] is None:
                        params = {
                            "fields1": "f1,f2,f3,f4,f5,f6,f7,f8,f9,f10,f11,f12,f13",
                            "fields2": "f51,f52,f53,f54,f55,f56,f57,f58",
                            "ut": "fa5fd1943c7b386f172d6893dbfba10b",
                            "iscr": "0",
                            "ndays": "5",
                            "secid": f"47.{symbol}",
                            "_": "1623766962675",
                        }
            r = requests.get(url, params=params)
            data_json = r.json()
            temp_df = pd.DataFrame(
                [item.split(",") for item in data_json["data"]["trends"]]
            )
            temp_df.columns = [
                "时间",
                "开盘",
                "收盘",
                "最高",
                "最低",
                "成交量",
                "成交额",
                "最新价",
            ]
            temp_df.index = pd.to_datetime(temp_df["时间"], errors="coerce")
            temp_df = temp_df[start_date:end_date]
            temp_df.reset_index(drop=True, inplace=True)
            temp_df["开盘"] = pd.to_numeric(temp_df["开盘"], errors="coerce")
            temp_df["收盘"] = pd.to_numeric(temp_df["收盘"], errors="coerce")
            temp_df["最高"] = pd.to_numeric(temp_df["最高"], errors="coerce")
            temp_df["最低"] = pd.to_numeric(temp_df["最低"], errors="coerce")
            temp_df["成交量"] = pd.to_numeric(temp_df["成交量"], errors="coerce")
            temp_df["成交额"] = pd.to_numeric(temp_df["成交额"], errors="coerce")
            temp_df["最新价"] = pd.to_numeric(temp_df["最新价"], errors="coerce")
            temp_df["时间"] = pd.to_datetime(temp_df["时间"]).astype(str)
            return temp_df
        else:
            url = "http://push2his.eastmoney.com/api/qt/stock/kline/get"
            try:
                params = {
                    "secid": f"{code_id_dict[symbol]}.{symbol}",
                    "ut": "7eea3edcaed734bea9cbfc24409ed989",
                    "fields1": "f1,f2,f3,f4,f5,f6",
                    "fields2": "f51,f52,f53,f54,f55,f56,f57,f58,f59,f60,f61",
                    "klt": period,
                    "fqt": "1",
                    "beg": "0",
                    "end": "20500000",
                    "_": "1630930917857",
                }
            except:
                params = {
                    "secid": f"0.{symbol}",
                    "ut": "7eea3edcaed734bea9cbfc24409ed989",
                    "fields1": "f1,f2,f3,f4,f5,f6",
                    "fields2": "f51,f52,f53,f54,f55,f56,f57,f58,f59,f60,f61",
                    "klt": period,
                    "fqt": "1",
                    "beg": "0",
                    "end": "20500000",
                    "_": "1630930917857",
                }
                r = requests.get(url, params=params)
                data_json = r.json()
                if data_json["data"] is None:
                    params = {
                        "secid": f"1.{symbol}",
                        "ut": "7eea3edcaed734bea9cbfc24409ed989",
                        "fields1": "f1,f2,f3,f4,f5,f6",
                        "fields2": "f51,f52,f53,f54,f55,f56,f57,f58,f59,f60,f61",
                        "klt": period,
                        "fqt": "1",
                        "beg": "0",
                        "end": "20500000",
                        "_": "1630930917857",
                    }
                    r = requests.get(url, params=params)
                    data_json = r.json()
                    if data_json["data"] is None:
                        params = {
                            "secid": f"47.{symbol}",
                            "ut": "7eea3edcaed734bea9cbfc24409ed989",
                            "fields1": "f1,f2,f3,f4,f5,f6",
                            "fields2": "f51,f52,f53,f54,f55,f56,f57,f58,f59,f60,f61",
                            "klt": period,
                            "fqt": "1",
                            "beg": "0",
                            "end": "20500000",
                            "_": "1630930917857",
                        }
            r = requests.get(url, params=params)
            data_json = r.json()
            temp_df = pd.DataFrame(
                [item.split(",") for item in data_json["data"]["klines"]]
            )
            temp_df.columns = [
                "时间",
                "开盘",
                "收盘",
                "最高",
                "最低",
                "成交量",
                "成交额",
                "振幅",
                "涨跌幅",
                "涨跌额",
                "换手率",
            ]
            temp_df.index = pd.to_datetime(temp_df["时间"], errors="coerce")
            temp_df = temp_df[start_date:end_date]
            temp_df.reset_index(drop=True, inplace=True)
            temp_df["开盘"] = pd.to_numeric(temp_df["开盘"], errors="coerce")
            temp_df["收盘"] = pd.to_numeric(temp_df["收盘"], errors="coerce")
            temp_df["最高"] = pd.to_numeric(temp_df["最高"], errors="coerce")
            temp_df["最低"] = pd.to_numeric(temp_df["最低"], errors="coerce")
            temp_df["成交量"] = pd.to_numeric(temp_df["成交量"], errors="coerce")
            temp_df["成交额"] = pd.to_numeric(temp_df["成交额"], errors="coerce")
            temp_df["振幅"] = pd.to_numeric(temp_df["振幅"], errors="coerce")
            temp_df["涨跌幅"] = pd.to_numeric(temp_df["涨跌幅"], errors="coerce")
            temp_df["涨跌额"] = pd.to_numeric(temp_df["涨跌额"], errors="coerce")
            temp_df["换手率"] = pd.to_numeric(temp_df["换手率"], errors="coerce")
            temp_df["时间"] = pd.to_datetime(temp_df["时间"], errors="coerce").astype(str)
            temp_df = temp_df[
                [
                    "时间",
                    "开盘",
                    "收盘",
                    "最高",
                    "最低",
                    "涨跌幅",
                    "涨跌额",
                    "成交量",
                    "成交额",
                    "振幅",
                    "换手率",
                ]
            ]
            return temp_df


if __name__ == "__main__":
        models=stock_index_data()
        index_zh_a_hist_df = models.index_zh_a_hist(
            symbol="800000",
            period="daily",
            start_date="19700101",
            end_date="22220101",
        )
        print(index_zh_a_hist_df)

        index_zh_a_hist_min_em_df = models.index_zh_a_hist_min_em(
            symbol="000001",
            period="1",
            start_date="2023-12-11 09:30:00",
            end_date="2023-12-11 19:00:00",
        )
        print(index_zh_a_hist_min_em_df)
